from keras.layers import Conv2D, MaxPooling2D, Dropout,  Flatten, Dense,GRU, \
    Lambda,Permute,TimeDistributed,Bidirectional
from keras.models import Input,Model
from keras import backend as K
from config import Config


Height = Config.img_height
Width = Config.img_width
rnn_size = 64
n_str = Config.n_str + 1
n_color = Config.n_color + 1
n_len = Config.n_len
conv_shape = (None, 11, 512)


def ctc_lambda_func(args):
    y_pred, labels, input_length, label_length = args
    return K.ctc_batch_cost(labels, y_pred, input_length, label_length)


def nn_base(input_tensor):

    conv1_1 = Conv2D(32, 3, activation='relu', padding='same', kernel_initializer='he_normal')(input_tensor)
    conv1_2 = Conv2D(32, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv1_1)
    pool1 = MaxPooling2D((2, 2))(conv1_2)
    conv2_1 = Conv2D(64, 3, activation='relu', padding='same', kernel_initializer='he_normal')(pool1)
    conv2_2 = Conv2D(64, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv2_1)
    pool2 = MaxPooling2D((2, 2))(conv2_2)
    conv3_1 = Conv2D(128, 3, activation='relu', padding='same', kernel_initializer='he_normal')(pool2)
    conv3_2 = Conv2D(128, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv3_1)
    pool3 = MaxPooling2D((2, 2))(conv3_2)

    m = Permute((2, 1, 3), name='permute')(pool3)

    flt = TimeDistributed(Flatten(), name='timedistrib')(m)

    des = Dense(32)(flt)

    gru_1 = Bidirectional(GRU(rnn_size, return_sequences=True, kernel_initializer='he_normal'), merge_mode='sum')(des)
    gru_2 = Bidirectional(GRU(rnn_size, return_sequences=True, kernel_initializer='he_normal'), merge_mode='concat')(
        gru_1)

    x = Dropout(0.25)(gru_2)
    x1 = Dense(n_str, kernel_initializer='he_normal', activation='softmax',name='str_output')(x)
    x2 = Dense(n_color, kernel_initializer='he_normal', activation='softmax', name='color_output')(x)

    return x1,x2


def ctc_model(input_tensor, return_layer):

    labels1 = Input(name='the_labels1', shape=[n_len], dtype='float32')
    input_length = Input(name='input_length1', shape=[1], dtype='int64')
    label_length = Input(name='label_length1', shape=[1], dtype='int64')
    loss_out1 = Lambda(ctc_lambda_func, output_shape=(1,),
                       name='ctc1')([return_layer, labels1, input_length, label_length])

    model = Model(inputs=[input_tensor, labels1, input_length, label_length], outputs=loss_out1)
    model.compile(loss={'ctc1': lambda y_true, y_pred: y_pred}, optimizer='adadelta')
    model.summary()
    return model


input_tensor = Input(shape=(Height, Width, 3))
return_layers = nn_base(input_tensor)

model_str = ctc_model(input_tensor, return_layers[0])
model_color = ctc_model(input_tensor, return_layers[1])
model_all = Model(input_tensor, [return_layers[0], return_layers[1]])
model_all.compile(loss='categorical_crossentropy', optimizer='adadelta', metrics=['accuracy'])



